# 基础镜像
FROM ubuntu:22.04

# 设置环境变量
ENV DEBIAN_FRONTEND=noninteractive \
    PATH=/root/miniconda3/bin:$PATH \
    # 设置构建参数
    ARCH="aarch64" \
    CANN_VERSION="8.0.0" \
    # 根据需要，可以是 "910b" 或 "310p"
    NPU_NAME="910b" \
    MIRRORS_ADDRESS=mirrors.huaweicloud.com \
    SRC_DIR=./create_image_files \
    MY_PACKAGEDIR=/tmp/package_dir 

# 拷贝当前目录下的安装包到容器中
RUN mkdir -p ${MY_PACKAGEDIR}
COPY ${SRC_DIR} ${MY_PACKAGEDIR}

# 使用 SHELL 指令切换到 Bash
SHELL ["/bin/bash", "-c"]

# 更换 apt 源和 pip 源，安装必要的工具并清理缓存
RUN \
    # 设置 apt 源为华为云镜像
    echo "deb http://mirrors.huaweicloud.com/ubuntu-ports/   jammy main restricted universe multiverse" > /etc/apt/sources.list && \
    echo "deb http://mirrors.huaweicloud.com/ubuntu-ports/   jammy-updates main restricted universe multiverse" >> /etc/apt/sources.list && \
    echo "deb http://mirrors.huaweicloud.com/ubuntu-ports/   jammy-security main restricted universe multiverse" >> /etc/apt/sources.list && \
    \
    # 更新 apt 并安装必要的工具
    apt update && \
    apt install --no-install-recommends -y gcc g++ make cmake zlib1g zlib1g-dev openssl libsqlite3-dev libssl-dev libffi-dev openssh-server && \
    apt install --no-install-recommends -y libbz2-dev libxslt1-dev unzip pciutils net-tools libblas-dev gfortran libblas3 ca-certificates && \
    apt install --no-install-recommends -y bc wget curl pkg-config liblapack3 vim liblapack-dev libhdf5-dev xz-utils cython3 python3-h5py sqlite3 && \
    apt install --no-install-recommends -y libopenblas-dev libgmpxx4ldbl liblzma-dev libxml2 libgl1-mesa-glx git numactl linux-tools-common ipmitool jq tzdata && \
    apt clean && \
    rm -rf /var/lib/apt/lists/* && \
    \
    # 设置时区为上海
    ln -snf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && echo "Asia/Shanghai" > /etc/timezone && \
    \
    # 安装 Miniconda
    echo "检查 ${MY_PACKAGEDIR} 下的文件：" && \
    ls -l ${MY_PACKAGEDIR} && \
    chmod +x ${MY_PACKAGEDIR}/Miniconda3*.sh && \
    ${MY_PACKAGEDIR}/Miniconda3*.sh -b -p /root/miniconda3 && \
    echo "source /root/miniconda3/bin/activate" >> /root/.bashrc && \
    \
    # 设置 pip 源为华为云镜像并配置超时
    mkdir -p /root/.pip && \
    echo -e "[global]\nindex-url = https://mirrors.huaweicloud.com/repository/pypi/simple\ntrusted-host   = mirrors.huaweicloud.com\ntimeout = 120" > /root/.pip/pip.conf && \
    \
    # 激活 Miniconda 环境并安装 Python 包
    source /root/miniconda3/bin/activate && \
    \
    # 清理 Miniconda 安装包
    rm -f ${MY_PACKAGEDIR}/Miniconda3*.sh && \
    \
    # 安装 Ascend 工具
    chmod +x ${MY_PACKAGEDIR}/Ascend*.run && \
    ${MY_PACKAGEDIR}/Ascend-cann-toolkit_${CANN_VERSION}_linux-${ARCH}.run --quiet --install && \
    echo 'export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver/:$LD_LIBRARY_PATH' >> /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    echo "source /usr/local/Ascend/ascend-toolkit/set_env.sh" >> /root/.bashrc && \
    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
    ${MY_PACKAGEDIR}/Ascend-cann-kernels-${NPU_NAME}_${CANN_VERSION}_linux-${ARCH}.run --quiet --install --install-for-all && \
    ${MY_PACKAGEDIR}/Ascend-cann-nnal_${CANN_VERSION}_linux-${ARCH}.run --quiet --install && \
    echo "source /usr/local/Ascend/nnal/atb/set_env.sh" >> /root/.bashrc && \
    \
    # 清理 Ascend 安装包
    rm -f ${MY_PACKAGEDIR}/Ascend-cann-*.run && \
    \
    # 安装 pip 包
    pip install --no-cache-dir "decorator~=5.1.1" \
        "attrs~=24.2.0" \
        "psutil~=6.0.0" \
        "absl-py~=2.1.0"  \
        "cloudpickle~=3.0.0"  \
        "scipy~=1.13.1"  \
        "synr~=0.5.0"  \
        "tornado~=6.4.1"  \
        "numpy~=1.26.4"  \
        "transformers~=4.37.0"  \
        "pandas~=2.2.2"  \
        "sentencepiece~=0.2.0"  \
        "accelerate~=0.34.2"  \
        "tiktoken~=0.7.0"  \
        "tensorboard~=2.17.1"  \
        "modelscope"  \
        "ipykernel"  \
        "torch==2.1.0" && \
    \
    # 安装 PyTorch NPU 版本
    mkdir -p ${MY_PACKAGEDIR}/pta && \
    tar -zxf ${MY_PACKAGEDIR}/pytorch_v2.1.0-6.0.rc3_py310.tar.gz -C ${MY_PACKAGEDIR}/pta && \
    pip install --no-cache-dir ${MY_PACKAGEDIR}/pta/torch_npu-2.1.0*${ARCH}.whl && \
    pip install --no-cache-dir ${MY_PACKAGEDIR}/pta/apex*${ARCH}.whl && \
    # 在另外一个虚拟环境中创建mindie 环境，基于base clone
    /root/miniconda3/bin/conda create -n mindie --clone base -y && \
    # 安装mindie
    bash ${MY_PACKAGEDIR}/Ascend-mindie_1.0.0_linux-${ARCH}.run --quiet --install && \
    echo "source /usr/local/Ascend/mindie/set_env.sh" >> /root/.bashrc && \
    mkdir -p /root/mindie && \
    tar -zxf ${MY_PACKAGEDIR}/Ascend-mindie-atb-models*_linux-${ARCH}_torch2.1.0-abi0.tar.gz -C /root/mindie && \
    echo "source /root/mindie/set_env.sh" >> /root/.bashrc && \
    # 清理安装包
    rm -rf ${MY_PACKAGEDIR}

